package com.practice.designMode.flyweight;

import java.nio.ByteBuffer;
import java.util.ArrayDeque;
import java.util.Deque;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * kafka的buffer pool
 */
public class BufferPool {
    private final int totalSize;
    private final int slotSize;
    private int free;
    // 享元模式，对于经常申请的长度对象进行一个封装复用，避免频繁创建
    private final Deque<ByteBuffer> slot = new ArrayDeque<>();
    private final Deque<Condition> waiters = new ArrayDeque<>();
    private final Lock lock = new ReentrantLock();

    public BufferPool(int totalSize, int slotSize) {
        this.totalSize = totalSize;
        this.slotSize = slotSize;
        this.free = totalSize;
    }

    public ByteBuffer allocate(int size, long timeout) throws InterruptedException {
        if(size > totalSize || size <= 0){
            throw new IllegalArgumentException("size is not valid");
        }
        lock.lock();
        try {
            if(size == slotSize && !slot.isEmpty()){
                return slot.pollFirst();
            }
            if((free + slot.size() * size) >= size){
                freeUp(size);
                free -= size;
                return ByteBuffer.allocate(size);
            }
            Condition condition = lock.newCondition();
            waiters.add(condition);
            long remainTime = timeout;
            try {
                while (true){
                    long start = System.currentTimeMillis();
                    Boolean wakeup = condition.await(remainTime, TimeUnit.MICROSECONDS);
                    if(!wakeup){
                        throw new InterruptedException("wait timeout");
                    }
                    if(size == slotSize && !slot.isEmpty()){
                        return slot.pollFirst();
                    }
                    if((free + slot.size() * size) >= size){
                        freeUp(size);
                        free -= size;
                        return ByteBuffer.allocate(size);
                    }
                    remainTime -= System.currentTimeMillis() - start;
                }
            }finally {
                waiters.remove(condition);
            }
        }finally {
            if(!waiters.isEmpty() && !(free == 0 && slot.isEmpty())){
                waiters.peekFirst().signal();
            }
            lock.unlock();
        }
    }

    private void freeUp(int size) {
        while (free < size && !slot.isEmpty()){
            free += slot.pollFirst().capacity();
        }
    }

    public void deallocate(ByteBuffer buffer){
        lock.lock();
        try{
            if(buffer.capacity() == slotSize){
                buffer.clear();
                slot.addLast(buffer);
            }else{
                free += buffer.capacity();
            }
            if(!waiters.isEmpty()){
                waiters.peekFirst().signal();
            }
        }finally {

        }
    }
}
